#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
编排器组件，负责协调多智能体之间的消息传递和流程管理
支持通过 MemorySystem 管理状态，并处理 Web 应用的交互请求
"""

from email import message
import logging
import uuid
import json
import re
import os
from datetime import datetime, timezone  # Make sure timezone is imported
from typing import Dict, List, Any, Optional, Tuple

from config.settings import SETTINGS
from core.persona import PersonaRepository

try:
    from agents.base_agent import BaseAgent
    from agents.companion import CompanionAgent
    from utils.memory_system import MemorySystem
    from config.settings import Settings
    from utils.llm_service import LLMService
except ImportError as e:
    logging.critical(f"Orchestrator 无法导入核心依赖: {e}", exc_info=True)

    class PlaceholderAgent:
        pass

    class PlaceholderMemorySystem:
        pass

    class PlaceholderLLMService:
        pass

    BaseAgent = PlaceholderAgent  # type: ignore
    CompanionAgent = PlaceholderAgent  # type: ignore
    MemorySystem = PlaceholderMemorySystem  # type: ignore
    LLMService = PlaceholderLLMService  # type: ignore
    STAGE_INTENT = STAGE_ASKING_IDENTITY = STAGE_GUIDING = STAGE_CHATTING = (
        STAGE_CHAT_ENDED
    ) = STAGE_CLARIFICATION = STAGE_ERROR = STAGE_COMPLETED = "unknown_stage"

logger = logging.getLogger("AI-MindCare-System-Orchestrator")


class Orchestrator:
    """编排器，负责协调多智能体系统的交互流程"""

    def __init__(
        self,
        memory_system: Optional[MemorySystem] = None,
        llm_service: Optional[LLMService] = None,
    ):
        if not isinstance(memory_system, MemorySystem):
            raise TypeError("必须提供有效的 MemorySystem 实例。")

        self.memory_system = memory_system
        self.llm_service = llm_service
        self.persona_repository = PersonaRepository()

        self.agents: Dict[str, BaseAgent] = {}

        # 初始化所有人设的CompanionAgent
        # for persona in SETTINGS.AGENT_PERSONAS.keys():
        for persona in self.persona_repository.list_available_personae():
            companion = CompanionAgent(
                persona=persona, memory_system=memory_system, llm_service=llm_service
            )

            self.agents[persona] = companion

        logger.info("编排器初始化完成，已完成所有人设的初始化。")

    def start_new_web_conversation(self, persona="professional") -> Tuple[str, str]:
        logger.info(f"编排器开始新会话，请求人设: {persona}")

        # 检查人设是否存在
        if persona not in self.agents:
            logger.warning(f"人设 {persona} 不存在，使用默认人设 professional")
            persona = "professional"

        agent = self.agents.get(persona)
        if not agent:
            logger.error(f"无法获取人设 {persona} 的agent")
            initial_message = "您好，我是您的 AI 助手。您可以告诉我具体需求，比如信息查询、任务协助或对话交流，我会尽力为您提供支持。"
        else:
            initial_message = agent.get_greetings()
            logger.info(f"从agent {persona} 获取的问候语: {initial_message}")

        if not initial_message:
            logger.warning(f"agent {persona} 返回的问候语为空，使用默认问候语")
            initial_message = "您好，我是您的 AI 助手。您可以告诉我具体需求，比如信息查询、任务协助或对话交流，我会尽力为您提供支持。"

        conversation_id = f"conv_{uuid.uuid4()}"  # 系统生成对话ID
        logger.info(f"生成会话ID: {conversation_id}")
        return conversation_id, initial_message

    def handle_web_message(
        self,
        conversation_id: str,
        user_message: str,  # 用户最新的消息
        user_buffer: Dict[str, Any],
        long_term_memory,
        persona="professional",
        model_key="deepseek-ai/DeepSeek-V3",
        image_base64=None,
        stream: bool = True,
    ) -> Dict[str, Any]:

        logger.info(
            f"Orchestrator: 处理 Web 消息 (Conversation: {conversation_id}): '{user_message[:50]}...'"
        )

        if not self.agents or not self.memory_system or not self.llm_service:
            logger.error(f"Orchestrator 缺少核心组件 (Conversation: {conversation_id})")
 
            return {
                "status": "error",
                "message": "系统内部错误，核心服务不可用。",
                "agent_type": "system",
            }
        target_agent = self.agents.get(persona)
        try:
            agent_message_envelope = {
                "conversation_id": conversation_id,
                "user_message": user_message,
                "user_buffer": user_buffer,
                "long_term_memory": long_term_memory,
                "persona": persona,
                "model_key": model_key,
                "image_base64": image_base64,  # 图像数据（如果有）
                "stream": stream,  # 是否使用流式响应
            }

            agent_response_msg = target_agent.receive_message(agent_message_envelope)

            return {
                "status": "success",
                "response": agent_response_msg.get("response"),
                "temperature": agent_response_msg.get("temperature"),
                "agent_type": "companion",
                "conversation_id": conversation_id,
                "stream": stream,
            }

        except Exception as e:
            logger.error(
                f"与 Agent: {target_agent.get_persona_info().get('name')}交互时出错 (Conversation: {conversation_id}): {e}",
                exc_info=True,
            )
            return {
                "status": "error",
                "message": f"系统在与 Agent: {target_agent.get_persona_info().get('name')}沟通时出现问题，请稍后重试。",

                "agent_type": "system",
            }
        
    def handle_web_message_stream(
        self,
        conversation_id: str,
        user_message: str,  # 用户最新的消息
        user_buffer: Dict[str, Any],
        long_term_memory,
        persona="professional",
        model_key="deepseek-ai/DeepSeek-V3",
        image_base64=None,
        stream: bool = True,
    ) -> Dict[str, Any]:

        logger.info(
            f"Orchestrator: 处理 Web 消息 (Conversation: {conversation_id}): '{user_message[:50]}...'"
        )

        if not self.agents or not self.memory_system or not self.llm_service:
            logger.error(f"Orchestrator 缺少核心组件 (Conversation: {conversation_id})")
 
            return {
                "status": "error",
                "message": "系统内部错误，核心服务不可用。",
                "agent_type": "system",
            }
        target_agent = self.agents.get(persona)
        try:
            agent_message_envelope = {
                "conversation_id": conversation_id,
                "user_message": user_message,
                "user_buffer": user_buffer,
                "long_term_memory": long_term_memory,
                "persona": persona,
                "model_key": model_key,
                "image_base64": image_base64,  # 图像数据（如果有）
                "stream": stream,  # 是否使用流式响应
            }
            
            for agent_response_msg in target_agent.receive_message_stream(agent_message_envelope):

                yield {
                    "status": "success",
                    "response": agent_response_msg.get("response"),
                    "temperature": agent_response_msg.get("temperature"),
                    "agent_type": "companion",
                    "conversation_id": conversation_id,
                    "stream": stream,
                }

        except Exception as e:
            logger.error(
                f"与 Agent: {target_agent.get_persona_info().get('name')}交互时出错 (Conversation: {conversation_id}): {e}",
                exc_info=True,
            )
            yield {
                "status": "error",
                "message": f"系统在与 Agent: {target_agent.get_persona_info().get('name')}沟通时出现问题，请稍后重试。",

                "agent_type": "system",
            }

    # --- END REINSTATED HELPER METHODS ---
